import torch

from eval import show_without_plt

import csv

def check_orthogonality(matrix):
    if matrix.size(0) != matrix.size(1):
        raise ValueError("输入矩阵必须是正方形矩阵")

    n = matrix.size(0)
    identity = torch.eye(n, dtype=matrix.dtype, device=matrix.device)
    product = torch.matmul(matrix, matrix.t())  # 计算矩阵乘以其转置
    error = torch.norm(product - identity)  # Frobenius 范数计算误差

    return error.item()

def det(matrix):
    return torch.det(matrix)

def rotation_matrix_eval(orthogonal_matrix):
    print("--------------------")
    print(f"旋转矩阵形状: {orthogonal_matrix.shape}")
    error = check_orthogonality(orthogonal_matrix)
    print(f"正交性误差: {error}")
    print(f"行列式: {det(orthogonal_matrix)}")
    print('-------------------')

def load_ffn_fm(path, range_limit=32, specific=False):
    if specific:
        return torch.stack([torch.load(f'{path}/ffn_feature_map_{i}.pth', weights_only=False).to(torch.float32) for i in range(0, range_limit)]).to(torch.device('cuda'))
    else:
        return torch.stack([torch.load(f'{path}/tensor{i}.pth', weights_only=False).to(torch.float32) for i in range(0, range_limit)]).to(torch.device('cuda'))

class Stats:
    def __init__(self, name=None):
        self.lower_radio = []
        self.lower_nums = []
        self.outliers_radio = []
        self.shape = None
        self.name = name
        # 稀疏度
        self.sparse = None
        # 异常值数量
        self.outliers = None
        # 异常值占比
        self.outliers_ratio = None
    def cal (self):
        self.sparse = sum(self.lower_radio) / len(self.lower_radio)
        self.outliers = sum(self.lower_nums) // 32
        self.outliers_ratio = sum(self.outliers_radio) / len(self.outliers_radio)


    def show (self):
        self.cal()
        print("--------------------")
        print(f"名称: {self.name}")
        print(f"矩阵形状: {self.shape}")
        print(f"1e-4以下的数量占比(稀疏度): {self.sparse}")
        print(f"1e-4以下的值的数量: {self.outliers}")
        print(f"异常值数量占比: {self.outliers_ratio}")
        print("--------------------")

    def save_csv(self, filename='dataStats.csv'):
        # 使用 'a' 模式追加数据，而不是覆盖
        with open(filename, 'a', newline='') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=["名称", "矩阵形状", "1e-4以下的数量占比(稀疏度)", "1e-4以下的值的数量", "异常值占比"])

            # 如果文件为空，写入表头
            if csvfile.tell() == 0:
                writer.writeheader()

            writer.writerow({
                "名称": self.name,
                "矩阵形状": self.shape,
                "1e-4以下的数量占比(稀疏度)": self.sparse,
                "1e-4以下的值的数量": self.outliers,
                "异常值占比": self.outliers_ratio
            })

# 处理将旋转后矩阵装入Stats类
def handle_stats(stats, matrix):
    for i in range(matrix.shape[0]):
        temp = show_without_plt(matrix[i])
        stats.lower_radio.append(temp[-1])
        stats.lower_nums.append(temp[2])
        stats.outliers_radio.append(temp[3]["outlier_ratio"])
    stats.shape = matrix.shape[1:]

if __name__ == "__main__":
    # 正交矩阵示例
    orthogonal_matrix = torch.load('trained_rotation_matrix.pth',weights_only=True)
    range_limit = 32

    first_line_input_fms = load_ffn_fm("ffn/first_line_input",range_limit)
    first_line_output_fms = load_ffn_fm("ffn/first_line_output",range_limit)
    second_line_input_fms_bef = load_ffn_fm('ffn/second_line_input', range_limit)
    second_line_output_fms_bef = load_ffn_fm('ffn/second_line_output', range_limit)
    second_line_output_fms_aft = load_ffn_fm('ffn/aft_second_line_output', range_limit)

    rotation_matrix_eval(orthogonal_matrix)

    first_line_input = Stats(name="第一线性层输入（FFN层的输入）")
    first_line_output = Stats(name="第一线性层输出（非线性层的输入）")
    second_line_input_bef = Stats(name="第二线性层输入（非线性层的输出）旋转前")
    second_line_input_aft = Stats(name="第二线性层输入（非线性层的输出）旋转后")
    second_line_output_bef = Stats(name="第二线性层输出（FFN层的输出）旋转前")
    second_line_output_aft = Stats(name="第二线性层输出（FFN层的输出）旋转后")

    stats = [first_line_input, first_line_output, second_line_input_bef, second_line_input_aft, second_line_output_bef, second_line_output_aft]

    handle_stats(first_line_input, second_line_input_fms_bef)
    handle_stats(first_line_output, second_line_output_fms_bef)
    handle_stats(second_line_input_bef, second_line_input_fms_bef)
    handle_stats(second_line_input_aft, torch.matmul(second_line_input_fms_bef, orthogonal_matrix))
    handle_stats(second_line_output_bef, second_line_output_fms_bef)
    handle_stats(second_line_output_aft, second_line_output_fms_aft)




    for stat in stats:
        stat.show()
        stat.save_csv()








